#!/usr/bin/python3
#coding:utf8

import os, sys, time, json, glob
from PIL import Image
import numpy as np
import cv2
import serial
import re
import argparse

param = {}

def get_dev(ver):
    path = 'None'
    video = 'None'
    ttyUSBs = glob.glob('/dev/ttyUSB*')
    #查找led灯
    try:
        paths = []
        for dev in ttyUSBs:
            ser = serial.Serial(dev, 9600)
            if ser is None:
                print("Fail to open dev: %s" % dev)
                continue
            ser.parity = 'N'
            ser.stopbits = 1
            ser.bytesize = 8
            if ver == 0:
                # FF 26 00 00 EE, 查询继电器状态
                ser.write(bytes.fromhex('ff260000ee'))
                time.sleep(0.02)
                ret = ser.read()
                if len(ret)>0:
                    paths.append(dev)
                    break
            elif ver == 1:
                # 30, 查询继电器状态
                ser.write(bytes.fromhex('30'))
                time.sleep(0.02)
                ret = ser.read()
                if len(ret)>0:
                    paths.append(dev)
                    break
            ser.close()
        path = paths[0]
    except Exception as e:
        print(e)

    #查找摄像头
    videos = glob.glob('/dev/video*')
    videos.reverse()
    videos = [v for i, v in enumerate(videos) if i % 2 == 0]
    if len(videos) >= 0:
        video = videos[0]

    print('relay: {}\ncamera: {}'.format(path, video))
    return path, video

def get_time_str():
    return time.strftime("%Y%m%d%H%M%S", time.localtime())

def print_info(cap):
    infos = [
       ('pos_msec', cv2.CAP_PROP_POS_MSEC),
       ('pos_frames', cv2.CAP_PROP_POS_FRAMES),
       ('pos_avi_ratio', cv2.CAP_PROP_POS_AVI_RATIO),
       ('frame_width', cv2.CAP_PROP_FRAME_WIDTH),
       ('frame_height', cv2.CAP_PROP_FRAME_HEIGHT),
       ('fps', cv2.CAP_PROP_FPS),
       ('fourcc', cv2.CAP_PROP_FOURCC),
       ('frame_count', cv2.CAP_PROP_FRAME_COUNT),
       ('format', cv2.CAP_PROP_FORMAT),
       ('mode', cv2.CAP_PROP_MODE),
       ('brightness', cv2.CAP_PROP_BRIGHTNESS),
       ('contrast', cv2.CAP_PROP_CONTRAST),
       ('saturation', cv2.CAP_PROP_SATURATION),
       ('hue', cv2.CAP_PROP_HUE),
       ('gain', cv2.CAP_PROP_GAIN),
       ('exposure', cv2.CAP_PROP_EXPOSURE),
       ('convert_rgb', cv2.CAP_PROP_CONVERT_RGB),
       ('white_balance_blue_u', cv2.CAP_PROP_WHITE_BALANCE_BLUE_U),
       ('rectification', cv2.CAP_PROP_RECTIFICATION),
       ('monochrome', cv2.CAP_PROP_MONOCHROME),
       ('sharpness', cv2.CAP_PROP_SHARPNESS),
       ('auto_exposure', cv2.CAP_PROP_AUTO_EXPOSURE),
       ('gamma', cv2.CAP_PROP_GAMMA),
       ('temperature', cv2.CAP_PROP_TEMPERATURE),
       ('trigger', cv2.CAP_PROP_TRIGGER),
       ('trigger_delay', cv2.CAP_PROP_TRIGGER_DELAY),
       ('white_balance_red_v', cv2.CAP_PROP_WHITE_BALANCE_RED_V),
       ('zoom', cv2.CAP_PROP_ZOOM),
       ('focus', cv2.CAP_PROP_FOCUS),
       ('guid', cv2.CAP_PROP_GUID),
       ('iso_speed', cv2.CAP_PROP_ISO_SPEED),
       ('backlight', cv2.CAP_PROP_BACKLIGHT),
       ('pan', cv2.CAP_PROP_PAN),
       ('tilt', cv2.CAP_PROP_TILT),
       ('roll', cv2.CAP_PROP_ROLL),
       ('iris', cv2.CAP_PROP_IRIS),
       ('settings', cv2.CAP_PROP_SETTINGS),
       ('buffersize', cv2.CAP_PROP_BUFFERSIZE),
       ('autofocus', cv2.CAP_PROP_AUTOFOCUS),
       ('sar_num', cv2.CAP_PROP_SAR_NUM),
       ('sar_den', cv2.CAP_PROP_SAR_DEN),
       ('backend', cv2.CAP_PROP_BACKEND),
       ('channel', cv2.CAP_PROP_CHANNEL),
       ('auto_wb', cv2.CAP_PROP_AUTO_WB),
       ('wb_temperature', cv2.CAP_PROP_WB_TEMPERATURE)
    ]
    for d, k in infos:
        print('{}: {}'.format(d, cap.get(k)))

def white_balance(img):
    height = img.shape[0]
    width = img.shape[1]

    # 1.计算三通道灰度平均值
    imgB = img[:, :, 0]
    imgG = img[:, :, 1]
    imgR = img[:, :, 2]

    # 下述3行代码控制白平衡或者冷暖色调，下例中增加了b的分量，会生成冷色调的图像，
    # 如要实现白平衡，则把两个+10都去掉；如要生成暖色调，则增加r的分量即可。
    #bAve = cv2.mean(imgB)[0]
    #gAve = cv2.mean(imgG)[0] + 10
    #rAve = cv2.mean(imgR)[0] + 10
    bAve = cv2.mean(imgB)[0]
    gAve = cv2.mean(imgG)[0]
    rAve = cv2.mean(imgR)[0]
    aveGray = (int)(bAve + gAve + rAve) / 3

    # 2计算每个通道的增益系数
    bCoef = aveGray / bAve
    gCoef = aveGray / gAve
    rCoef = aveGray / rAve

    # 3使用增益系数
    imgB = np.floor((imgB * bCoef))  # 向下取整
    imgG = np.floor((imgG * gCoef))
    imgR = np.floor((imgR * rCoef))

    imgB = np.minimum(imgB, np.ones(imgB.shape)*255).astype(np.uint8)
    imgG = np.minimum(imgG, np.ones(imgB.shape)*255).astype(np.uint8)
    imgR = np.minimum(imgR, np.ones(imgB.shape)*255).astype(np.uint8)

    dst = np.zeros(img.shape, dtype=img.dtype)
    dst[:, :, 0] = imgB
    dst[:, :, 1] = imgG
    dst[:, :, 2] = imgR

    return dst

def crop():
    global param
    work_dir = "."

    img_label_pth = "{}/label.jpg".format(work_dir)
    if not os.path.isfile(img_label_pth):
        print("[%s] There is no image: %s" % (get_time_str(), img_label_pth))
        return False

    img_cntnt_pth = "{}/cntnt.jpg".format(work_dir)
    if not os.path.isfile(img_cntnt_pth):
        print("[%s] There is no image: %s" % (get_time_str(), img_cntnt_pth))
        return False

    res_w, res_h = param.get('res')
    img_label = Image.open(img_label_pth)
    if img_label.size[0] != res_w or img_label.size[1] != res_h:
        print("[%s] Invalid image size: %s, image: %s, Resolution: [%d, %d]" \
                % (get_time_str(), str(img_label.size), img_label_pth, nResolutionW, nResolutionH))
        return False

    img_cntnt = Image.open(img_cntnt_pth)
    if img_cntnt.size[0] != res_w or img_cntnt.size[1] != res_h:
        print("[%s] Invalid image size: %s, image: %s, Resolution: [%d, %d]" \
                % (get_time_str(), str(img_cntnt.size), img_cntnt_path, nResolutionW, nResolutionH))
        return False

    gls_img_w, gls_img_h = param.get('gls')
    gls1_x, gls1_y = param.get('gls1')
    gls2_x, gls2_y = param.get('gls2')
    gls3_x, gls3_y = param.get('gls3')
    gls4_x, gls4_y = param.get('gls4')
    label_gls1 = img_label.crop((gls1_x, gls1_y, gls1_x + gls_img_w, gls1_y + gls_img_h))
    label_gls2 = img_label.crop((gls2_x, gls2_y, gls2_x + gls_img_w, gls2_y + gls_img_h))
    label_gls3 = img_label.crop((gls3_x, gls3_y, gls3_x + gls_img_w, gls3_y + gls_img_h))
    label_gls4 = img_label.crop((gls4_x, gls4_y, gls4_x + gls_img_w, gls4_y + gls_img_h))

    cntnt_gls1 = img_cntnt.crop((gls1_x, gls1_y, gls1_x + gls_img_w, gls1_y + gls_img_h))
    cntnt_gls2 = img_cntnt.crop((gls2_x, gls2_y, gls2_x + gls_img_w, gls2_y + gls_img_h))
    cntnt_gls3 = img_cntnt.crop((gls3_x, gls3_y, gls3_x + gls_img_w, gls3_y + gls_img_h))
    cntnt_gls4 = img_cntnt.crop((gls4_x, gls4_y, gls4_x + gls_img_w, gls4_y + gls_img_h))

    label_gls1.save("{}/Gls1_label.jpg".format(work_dir))
    cntnt_gls1.save("{}/Gls1_cntnt.jpg".format(work_dir))
    label_gls2.save("{}/Gls2_label.jpg".format(work_dir))
    cntnt_gls2.save("{}/Gls2_cntnt.jpg".format(work_dir))
    label_gls3.save("{}/Gls3_label.jpg".format(work_dir))
    cntnt_gls3.save("{}/Gls3_cntnt.jpg".format(work_dir))
    label_gls4.save("{}/Gls4_label.jpg".format(work_dir))
    cntnt_gls4.save("{}/Gls4_cntnt.jpg".format(work_dir))

    if not os.path.isdir('picture'):
        os.system('mkdir picture')
    label_gls1.save("{}/picture/{}_Gls1_label.jpg".format(work_dir, time.strftime('%y%m%d%H%M%S', time.localtime())))
    cntnt_gls1.save("{}/picture/{}_Gls1_cntnt.jpg".format(work_dir, time.strftime('%y%m%d%H%M%S', time.localtime())))
    label_gls2.save("{}/picture/{}_Gls2_label.jpg".format(work_dir, time.strftime('%y%m%d%H%M%S', time.localtime())))
    cntnt_gls2.save("{}/picture/{}_Gls2_cntnt.jpg".format(work_dir, time.strftime('%y%m%d%H%M%S', time.localtime())))
    label_gls3.save("{}/picture/{}_Gls3_label.jpg".format(work_dir, time.strftime('%y%m%d%H%M%S', time.localtime())))
    cntnt_gls3.save("{}/picture/{}_Gls3_cntnt.jpg".format(work_dir, time.strftime('%y%m%d%H%M%S', time.localtime())))
    label_gls4.save("{}/picture/{}_Gls4_label.jpg".format(work_dir, time.strftime('%y%m%d%H%M%S', time.localtime())))
    cntnt_gls4.save("{}/picture/{}_Gls4_cntnt.jpg".format(work_dir, time.strftime('%y%m%d%H%M%S', time.localtime())))

    return True

def rotate(img, angle, center=None, scale=1.0):
    (h, w) = img.shape[:2]
    if center is None:
        center = (w//2, h//2)
    M = cv2.getRotationMatrix2D(center, angle, scale)
    rotated = cv2.warpAffine(img, M, (w, h))
    return rotated

def take_photo(is_real, wait_cnt, expo_time):
    global param
    frame = None
    # 打开摄像头
    camera = param.get('camera') 
    cam_idx = re.search(r'\d+', camera).group()
    #print(f'cam_idx: {cam_idx}')
    cap = cv2.VideoCapture(int(cam_idx)) # 此处填入摄像头
    if cap is None:
        print("[%s] Fail to open camera: %d" % (get_time_str(), camera))
        return frame

    #cap.set(cv2.CAP_PROP_AUTO_EXPOSURE, 0) # 切换为手动曝光
    cap.set(cv2.CAP_PROP_EXPOSURE, expo_time) # 设置曝光时间

    res_w, res_h = param.get('res')
    if res_w != cap.get(cv2.CAP_PROP_FRAME_WIDTH) or res_h != cap.get(cv2.CAP_PROP_FRAME_HEIGHT):
        #print(res_w, res_h)
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, res_w)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, res_h)

    #print_info(cap)

    if is_real:
        #err_cnt = 0
        while True:
            #if err_cnt > 5: break
            ret, frame = cap.read() # get frame
            if frame is None:
                #err_cnt += 1
                break
            else:
                cv2.imshow(f"real", frame)
                if cv2.waitKey(1)&0xFF == ord('q') or ret == False: break
        cv2.destroyAllWindows()
    else:
        #err_cnt = 0
        while True:
            #if err_cnt > 5: break
            ret, frame = cap.read() # get frame
            if frame is None:
                #err_cnt += 1
                break
            else:
                if wait_cnt == 0: break
                else: wait_cnt -= 1
        if frame is not None:
            frame = rotate(frame, param.get('rotate'))
            frame = white_balance(frame)

    cap.release()
    return frame

def snapshot(model):
    global param
    ret = True
    set_light('ALL', 'OFF')
    while True:
        if model == 'snapshot':
            # 仅点亮上灯
            set_light('TOP', 'ON')

            # 拍照
            frame = take_photo(False, param.get('label_wait_cnt'), param.get('label_expo_time'))
            if frame is None:
                print("[%s] take_photo label failed" % get_time_str())
                ret = False
                break
            cv2.imwrite('label.jpg', frame)

            # 仅点亮下灯
            set_light('BOTTOM', 'ON')

            # 拍照
            frame = take_photo(False, param.get('cntnt_wait_cnt'), param.get('cntnt_expo_time'))
            if frame is None:
                print("[%s] take_photo cntnt failed" % get_time_str())
                ret = False
                break
            cv2.imwrite('cntnt.jpg', frame)

            # 切图
            crop()
        elif model == 'labelreal':
            # 仅点亮上灯
            set_light('TOP', 'ON')

            # 显示实时图
            frame = take_photo(True, param.get('label_wait_cnt'), param.get('label_expo_time'))
            if frame is None:
                print("[%s] take_photo label failed" % get_time_str())
                ret = False
                break
            cv2.imwrite(f'{model}.jpg', frame)
        elif model == 'cntntreal':
            # 仅点亮下灯
            set_light('BOTTOM', 'ON')

            # 显示实时图
            frame = take_photo(True, param.get('cntnt_wait_cnt'), param.get('cntnt_expo_time'))
            if frame is None:
                print("[%s] take_photo cntnt failed" % get_time_str())
                ret = False
                break
            cv2.imwrite(f'{model}.jpg', frame)
        elif model == 'cameraonly':
            # 显示实时图
            frame = take_photo(True, 1, param.get('default_expo_time'))
            if frame is None:
                print("[%s] take_photo default failed" % get_time_str())
                ret = False
                break
            cv2.imwrite(f'{model}.jpg', frame)
        break
    set_light('ALL', 'OFF')
    return ret

def set_light_v0(pos, action):
    # 打开继电器
    if param.get('relay') != 'None':
        ser = serial.Serial(param.get('relay'), 9600)
        if ser is None:
            print('Open relay fail.')
            return
        ser.parity = 'N'
        ser.stopbits = 1
        ser.bytesize = 8

    if ser is not None:
        if action == 'ON':
            if pos == 'TOP':#上路
                ser.write(b'\xff\x02\x01\x03\xee')
            elif pos == 'BOTTOM':#下路
                ser.write(b'\xff\x25\x00\x25\xee')
                time.sleep(1)
                ser.write(b'\xff\x01\x01\x02\xee')
            else:
                print(f'Invalid Index: {pos}')
        elif action == 'OFF': 
            ser.write(b'\xff\x25\x00\x25\xee')
    ser.close()

def set_light_v1(pos, action):
    # 0, 为继电器CH1, 1, 为继电器CH2
    status = {
            'ON': ['0001', '0010'],
            'ONALL': '0011',
            'OFFALL': '0000',
            }

    # 打开继电器
    if param.get('relay') != 'None':
        ser = serial.Serial(param.get('relay'), 9600)
        if ser is None:
            print('Open relay fail.')
            return
        ser.parity = 'N'
        ser.stopbits = 1
        ser.bytesize = 8

    if ser is not None:
        cmd = '0000'
        if action == 'ON':
            if pos == 'TOP':
                cmd += status['ON'][1]
            elif pos == 'BOTTOM':
                cmd += status['ON'][0]
        elif action == 'OFF':
            cmd = '00101111'
            #cmd += status['OFFALL']
        else:
            print(f'Invalid Index: {pos}')

        cmd = '{:02x}'.format(int(cmd, 2))
        ser.write(bytes.fromhex(cmd))
        ret = ser.read()
    ser.close()

def set_light(pos, action):
    global param
    if param.get('relay') is not None:
        ver = param.get('devver')
        if ver == 0:
            set_light_v0(pos, action)
            if action != 'OFF': time.sleep(0.1)
        elif ver == 1:
            set_light_v1(pos, action)
            if action != 'OFF': time.sleep(0.1)
        else:
            print(f'Invalid DevVer: {ver}')

def set_param(cfg):
    global param
    if cfg is not None:
        param['camera'] = cfg.get('strCameraPath') 
        param['relay'] = cfg.get('strLightsPath') 
        param['label_expo_time'] = cfg.get('nExpTimeLabel') 
        param['cntnt_expo_time'] = cfg.get('nExpTimeCntnt') 
        param['label_wait_cnt'] = cfg.get('nWaitFrameCntLabel') 
        param['cntnt_wait_cnt'] = cfg.get('nWaitFrameCntCntnt') 
        param['devver'] = cfg.get('nDevVer') 
        param['res'] = (cfg.get('nResolutionW'), cfg.get('nResolutionH'))
        param['rotate'] = cfg.get('nImgRotate')
        param['gls'] = (cfg.get('nGlsImgW'), cfg.get('nGlsImgH'))
        param['gls1'] = (cfg.get('nGls1PosX'), cfg.get('nGls1PosY'))
        param['gls2'] = (cfg.get('nGls2PosX'), cfg.get('nGls2PosY'))
        param['gls3'] = (cfg.get('nGls3PosX'), cfg.get('nGls3PosY'))
        param['gls4'] = (cfg.get('nGls4PosX'), cfg.get('nGls4PosY'))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', type=str, default='snapshot', choices=['snapshot', 'labelreal', 'cntntreal', 'cameraonly'])
    parser.add_argument('--relay', type=str, default='None', help='relay device file, like /dev/ttyUSB*.')
    parser.add_argument('--camera', type=str, default='None', help='camera device file, like /dev/video*.')
    parser.add_argument('--devver', type=int, default=1, help='device version, [0, 1].')
    parser.add_argument('--cfg', type=str, default='./cfg.json', help='cfg.json file path.')
    args = parser.parse_args()

    if os.path.isfile(args.cfg):
        with open(args.cfg, 'r', encoding='utf-8') as f:
            cfg = json.load(f)
            set_param(cfg.get('Snapshot'))

    if param.get('camera') is None or param.get('relay') is None:
        #res = [(320, 240), (352, 288), (640, 480), (800, 600), (1024, 768), (1280, 720), (1280, 1024), (1920, 1080)]
        #res = [(640, 480), (2592, 1944)]
        relay, camera = get_dev(args.devver)
        if args.camera != 'None': camera = args.camera
        if args.relay != 'None': relay = args.relay

    if param.get('camera'           ) is None: param['camera'] = camera
    if param.get('relay'            ) is None: param['relay'] = relay
    if param.get('label_expo_time'  ) is None: param['label_expo_time'] = 22
    if param.get('cntnt_expo_time'  ) is None: param['cntnt_expo_time'] = 4
    if param.get('default_expo_time') is None: param['default_expo_time'] = 4
    if param.get('label_wait_cnt'   ) is None: param['label_wait_cnt'] = 1
    if param.get('cntnt_wait_cnt'   ) is None: param['cntnt_wait_cnt'] = 1
    if param.get('devver'           ) is None: param['devver'] = args.devver
    if param.get('res'              ) is None: param['res'] = (1280, 1024)
    if param.get('rotate'           ) is None: param['rotate'] = 180

    if snapshot(args.model) is False:
        exit(1)
    exit(0)
